import pynvml
import numpy as np


def choose_optimal_device(num_gpus):
    rest_memo_rate = []
    for gpu_id in range(num_gpus):
        pynvml.nvmlInit()
        meminfo = pynvml.nvmlDeviceGetMemoryInfo(
            pynvml.nvmlDeviceGetHandleByIndex(gpu_id))  # 指定显卡号
        rest_memo_rate.append(meminfo.free/meminfo.total*100)
        print(f'{gpu_id}号显卡当前剩余现存：{rest_memo_rate[gpu_id]:.2f}%')
    best_choice = np.argmax(rest_memo_rate)
    print(f'应选择{best_choice}号显卡')
